#!/usr/bin/env python3
"""
单个CN-CLIP模型完整数据集评估脚本
用于独立评估单个模型，避免全部模型串行执行时间过长的问题
"""

import os
import sys
import json
import torch
import logging
from datetime import datetime
from pathlib import Path
import argparse

# 添加项目路径
sys.path.append("/root/WenwuClip")
from codebase.dataset.dataset import get_test_set
from codebase.cnclip_raw.model_loader import CNClipModelLoader  
from codebase.cnclip_raw.evaluator import CNClipEvaluator
from codebase.cnclip_raw.wenwu_dataset_adapter import create_adapted_dataloader

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def evaluate_single_model(model_name, device=None, results_dir="research/experiments/baseline", 
                         batch_size=8, max_samples=None):
    """评估单个CN-CLIP模型在完整文物数据集上的性能"""
    
    # 设备配置
    if device is None:
        if torch.cuda.is_available():
            device = "cuda"
        elif torch.backends.mps.is_available():
            device = "mps"
        else:
            device = "cpu"
    
    logger.info(f"开始评估模型: {model_name}")
    logger.info(f"使用设备: {device}")
    logger.info(f"批大小: {batch_size}")
    logger.info(f"最大样本数: {max_samples or '全部'}")
    
    # 创建结果目录
    results_dir = Path(results_dir)
    results_dir.mkdir(parents=True, exist_ok=True)
    
    # 时间戳
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    try:
        # 加载模型
        logger.info("正在加载模型...")
        model_loader = CNClipModelLoader(device=device)
        if not model_loader.load_model(model_name, download_root="./pretrained_weights"):
            logger.error(f"模型 {model_name} 加载失败")
            return False
        
        # 准备数据集 - 使用标准化测试集函数
        logger.info("正在准备测试集数据...")
        dataset = get_test_set(data_scale=1.0)
        
        # 如果指定最大样本数，截取数据集
        if max_samples and max_samples < len(dataset):
            dataset.data = dataset.data[:max_samples]
        
        logger.info(f"实际评估样本数: {len(dataset)}")
        
        # 创建数据加载器
        dataloader = create_adapted_dataloader(
            dataset, 
            batch_size=batch_size, 
            shuffle=False,
            num_workers=2
        )
        
        # 初始化评估器
        evaluator = CNClipEvaluator(model_loader, device=device)
        
        # 执行评估
        logger.info("开始执行评估...")
        results = evaluator.evaluate_retrieval(
            dataloader,
            k_values=[1, 3, 5, 10, 20, 50],
            max_samples=max_samples,
            save_features=False  # 不保存特征以节省内存
        )
        
        # 添加评估配置信息
        results['evaluation_config'] = {
            'model_name': model_name,
            'device': device,
            'actual_samples': len(dataset),
            'batch_size': batch_size,
            'max_samples': max_samples,
            'timestamp': timestamp
        }
        
        # 生成报告
        report = evaluator.generate_report(results)
        
        # 保存结果
        model_results_path = results_dir / f"CN-CLIP_{model_name}_完整评估_结果.json"
        model_report_path = results_dir / f"CN-CLIP_{model_name}_完整评估_报告.md"
        
        with open(model_results_path, 'w', encoding='utf-8') as f:
            json.dump(results, f, indent=2, ensure_ascii=False, default=str)
        
        with open(model_report_path, 'w', encoding='utf-8') as f:
            f.write(report)
        
        # 打印报告
        print(f"\n{'='*80}")
        print(f"模型 {model_name} 评估完成！")
        print(f"{'='*80}")
        print(report)
        print(f"\n结果文件: {model_results_path}")
        print(f"报告文件: {model_report_path}")
        print(f"{'='*80}")
        
        logger.info(f"模型 {model_name} 评估成功完成")
        return True
        
    except Exception as e:
        logger.error(f"评估模型 {model_name} 时出错: {e}")
        import traceback
        logger.error(traceback.format_exc())
        return False
        
    finally:
        # 清理GPU内存
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


def main():
    parser = argparse.ArgumentParser(description="单个CN-CLIP模型完整数据集评估")
    parser.add_argument('--model', type=str, required=True,
                       choices=['ViT-B-16', 'ViT-L-14', 'ViT-L-14-336', 'ViT-H-14', 'RN50'],
                       help='要评估的模型名称')
    parser.add_argument('--device', type=str, default=None,
                       help='计算设备 (cuda/mps/cpu)')
    parser.add_argument('--batch-size', type=int, default=8,
                       help='批大小 (默认: 8)')
    parser.add_argument('--max-samples', type=int, default=None,
                       help='最大样本数限制 (默认: 无限制)')
    parser.add_argument('--results-dir', type=str, default="research/experiments/baseline",
                       help='结果保存目录')
    
    args = parser.parse_args()
    
    # 执行评估
    success = evaluate_single_model(
        model_name=args.model,
        device=args.device,
        results_dir=args.results_dir,
        batch_size=args.batch_size,
        max_samples=args.max_samples
    )
    
    if success:
        print(f"✅ 模型 {args.model} 评估成功完成！")
    else:
        print(f"❌ 模型 {args.model} 评估失败！")
        sys.exit(1)


if __name__ == "__main__":
    main()